#ifndef YUV420PWIDGET_H
#define YUV420PWIDGET_H

#include <QOpenGLWidget>
#include <QOpenGLShaderProgram>
#include <QOpenGLFunctions>

class YUV420PWidget : public QOpenGLWidget, protected QOpenGLFunctions
{
public:
    YUV420PWidget(QWidget *parent = nullptr) : QOpenGLWidget(parent) {}

    void updateFrame(const uint8_t* yData, const uint8_t* uData, const uint8_t* vData,
                    int width, int height) {
        m_width = width;
        m_height = height;
        m_yData = yData;
        m_uData = uData;
        m_vData = vData;
        update();
    }

protected:
    void initializeGL() override {
        initializeOpenGLFunctions();

        // 编译着色器
        m_program.addShaderFromSourceCode(QOpenGLShader::Vertex,
            "#version 330 core\n"
            "layout(location = 0) in vec4 vertexIn;\n"
            "layout(location = 1) in vec2 textureIn;\n"
            "out vec2 textureCoord;\n"
            "void main() {\n"
            "    gl_Position = vertexIn;\n"
            "    textureCoord = textureIn;\n"
            "}");

        m_program.addShaderFromSourceCode(QOpenGLShader::Fragment,
            "#version 330 core\n"
            "in vec2 textureCoord;\n"
            "out vec4 fragColor;\n"
            "uniform sampler2D tex_y;\n"
            "uniform sampler2D tex_u;\n"
            "uniform sampler2D tex_v;\n"
            "const mat3 yuv2rgb = mat3(\n"
            "    1.164383, 1.164383, 1.164383,\n"        // Y系数
            "    0.0,     -0.391762, 2.017232,\n"         // U系数
            "    -0.812968, 0.0,      0.0);\n"            // V系数
            "void main() {\n"
            "    float y = texture(tex_y, textureCoord).r;\n"
            "    float u = texture(tex_u, textureCoord).r;\n"
            "    float v = texture(tex_v, textureCoord).r;\n"
            "    y = 1.1643 * (y - 0.0625);\n"
            "    u = u - 0.5;\n"
            "    v = v - 0.5;\n"
            "    fragColor = vec4(\n"
            "        y + 2.017232 * v,\n"
            "        y - 0.391762 * u - 0.812968 * v,\n"
            "        y + 1.164383 * u,\n"
            "        1.0\n"
            "    );\n"
            "}");
        m_program.link();

        // 初始化三个纹理
        glGenTextures(1, &m_yTexture);
        glGenTextures(1, &m_uTexture);
        glGenTextures(1, &m_vTexture);

        // 指定颜色缓冲区的清除值(背景色)
        glClearColor(0.0f, 0.0f, 0.0f, 1.0f);
    }

    void resizeGL(int w, int h) override {
#if 1
        //计算需要显示图片的窗口大小，用于实现长宽等比自适应显示
        if(m_size.width()  < 0 || m_size.height() < 0) return;

        if((double(w) / h) < (double(m_size.width()) / m_size.height()))
        {
            m_zoomSize.setWidth(w);
            m_zoomSize.setHeight(((double(w) / m_size.width()) * m_size.height()));
        }
        else
        {
            m_zoomSize.setHeight(h);
            m_zoomSize.setWidth((double(h) / m_size.height()) * m_size.width());
        }
        m_pos.setX(double(w - m_zoomSize.width()) / 2);
        m_pos.setY(double(h - m_zoomSize.height()) / 2);
        this->update(QRect(0, 0, w, h));
#else
       glViewport(0, 0, w, h);
#endif
    }

    void paintGL() override {
        if (!m_yData || !m_uData || !m_vData) return;

        glClearColor(0.0f, 0.0f, 0.0f, 1.0f);
        glClear(GL_COLOR_BUFFER_BIT);

#if 1
        glViewport(m_pos.x(), m_pos.y(), m_zoomSize.width(), m_zoomSize.height());  // 设置视图大小实现图片自适应
#endif

        m_program.bind();

        // 更新Y纹理
        glActiveTexture(GL_TEXTURE0);
        glBindTexture(GL_TEXTURE_2D, m_yTexture);
        glTexImage2D(GL_TEXTURE_2D, 0, GL_RED, m_width, m_height, 0,
                    GL_RED, GL_UNSIGNED_BYTE, m_yData);
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR);
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR);
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE);
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE);

        // 更新U纹理
        glActiveTexture(GL_TEXTURE1);
        glBindTexture(GL_TEXTURE_2D, m_uTexture);
        glTexImage2D(GL_TEXTURE_2D, 0, GL_RED, m_width/2, m_height/2, 0,
                    GL_RED, GL_UNSIGNED_BYTE, m_uData);
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR);
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR);
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE);
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE);

        // 更新V纹理
        glActiveTexture(GL_TEXTURE2);
        glBindTexture(GL_TEXTURE_2D, m_vTexture);
        glTexImage2D(GL_TEXTURE_2D, 0, GL_RED, m_width/2, m_height/2, 0,
                    GL_RED, GL_UNSIGNED_BYTE, m_vData);
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR);
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR);
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE);
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE);

        // 设置纹理单元
        m_program.setUniformValue("tex_y", 0);
        m_program.setUniformValue("tex_u", 1);
        m_program.setUniformValue("tex_v", 2);

        // 绘制矩形
        static const GLfloat vertices[] = {
            -1.0f, -1.0f,
             1.0f, -1.0f,
            -1.0f, 1.0f,
             1.0f, 1.0f,
        };
        static const GLfloat texCoords[] = {
            0.0f, 1.0f,
            1.0f, 1.0f,
            0.0f, 0.0f,
            1.0f, 0.0f,
        };

        glVertexAttribPointer(0, 2, GL_FLOAT, GL_FALSE, 0, vertices);
        glVertexAttribPointer(1, 2, GL_FLOAT, GL_FALSE, 0, texCoords);
        glEnableVertexAttribArray(0);
        glEnableVertexAttribArray(1);

        glDrawArrays(GL_TRIANGLE_STRIP, 0, 4);

        m_program.release();
    }

private:
    QOpenGLShaderProgram m_program;
    GLuint m_yTexture = 0;
    GLuint m_uTexture = 0;
    GLuint m_vTexture = 0;
    const uint8_t* m_yData = nullptr;
    const uint8_t* m_uData = nullptr;
    const uint8_t* m_vData = nullptr;
    int m_width = 0;
    int m_height = 0;
    QSize  m_size;
    QSizeF  m_zoomSize;
    QPointF m_pos;
};


#endif // YUV420PWIDGET_H
